{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "61c9418b-651a-4c91-80f1-39aed03f8627",
   "metadata": {},
   "outputs": [],
   "source": [
    "import openai\n",
    "import re\n",
    "import json\n",
    "import spacy\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3e8d918e-3aa4-442a-b205-71bbdae95f6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "nlp = spacy.load(\"en_core_web_sm\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "3ab1a10f-7c90-42e8-b1e7-6bd534c12b74",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = json.load(open('MathQA/train.json'))\n",
    "dev_data = json.load(open('MathQA/dev.json'))\n",
    "test_data = json.load(open('MathQA/test.json'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "7aff75db-7a4d-4862-bdf8-5754f461dcd5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# test_idx = np.random.choice(len(test_data), 600, replace=False)\n",
    "# np.save('MathQA/test_idx.npy', test_idx)\n",
    "\n",
    "test_idx = np.load('MathQA/test_idx.npy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "52c44a5e-c603-45ef-bb97-4d470377bd25",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_data = [test_data[i] for i in test_idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e022cfef-e94c-48ad-8e4a-abefb52b8818",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_original = open('MathQA/prompt_original.txt').read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "42e36d54-2e4e-4350-8ee1-e205920e13db",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Question: There are 15 trees in the grove. Grove workers will plant trees in the grove today. After they are done, there will be 21 trees. How many trees did the grove workers plant today?\n",
      "Options: a ) 3 , b ) 5 , c ) 6 , d ) 10 , e ) none of these\n",
      "Let's think step by step\n",
      "There are 15 trees originally.\n",
      "Then there were 21 trees after some more were planted.\n",
      "So there must have been 21 - 15 = 6.\n",
      "The answer is c, 6.\n",
      "\n",
      "Question: If there are 3 cars in the parking lot and 2 more cars arrive, how many cars are in the parking lot?\n",
      "Let's think step by step\n",
      "Options: a ) 3 , b ) 5 , c ) 6 , d ) 10 , e ) none of these\n",
      "There are originally 3 cars.\n",
      "2 more cars arrive.\n",
      "3 + 2 = 5.\n",
      "The answer is b, 5.\n",
      "\n",
      "Question: Leah had 32 chocolates and her sister had 42. If they ate 35, how many pieces do they have left in total?\n",
      "Options: a ) 39 , b ) 5 , c ) 28 , d ) 37 , e ) 33\n",
      "Let's think step by step\n",
      "Originally, Leah had 32 chocolates.\n",
      "Her sister had 42.\n",
      "So in total they had 32 + 42 = 74.\n",
      "After eating 35, they had 74 - 35 = 39.\n",
      "The answer is a, 39.\n",
      "\n",
      "Question: Jason had 20 lollipops. He gave Denny some lollipops. Now Jason has 12 lollipops. How many lollipops did Jason give to Denny?\n",
      "Options: a ) 3 , b ) 5 , c ) 6 , d ) 8 , e ) none of these\n",
      "Let's think step by step\n",
      "Jason started with 20 lollipops.\n",
      "Then he had 12 after giving some to Denny.\n",
      "So he gave Denny 20 - 12 = 8.\n",
      "The answer is d, 8.\n",
      "\n",
      "Question: Shawn has five toys. For Christmas, he got two toys each from his mom and dad. How many toys does he have now?\n",
      "Options: a ) 3 , b ) 5 , c ) 6 , d ) 9 , e ) none of these\n",
      "Let's think step by step\n",
      "Shawn started with 5 toys.\n",
      "If he got 2 toys each from his mom and dad, then that is 4 more toys.\n",
      "5 + 4 = 9.\n",
      "The answer is d, 9.\n",
      "\n",
      "Question: There were nine computers in the server room. Five more computers were installed each day, from monday to thursday. How many computers are now in the server room?\n",
      "Options: a ) 31 , b ) 29 , c ) 67 , d ) 18 , e ) none of these\n",
      "Let's think step by step\n",
      "There were originally 9 computers.\n",
      "For each of 4 days, 5 more computers were added.\n",
      "So 5 * 4 = 20 computers were added.\n",
      "9 + 20 is 29.\n",
      "The answer is b, 29.\n",
      "\n",
      "Question: Michael had 58 golf balls. On tuesday, he lost 23 golf balls. On wednesday, he lost 2 more. How many golf balls did he have at the end of wednesday?\n",
      "Options: a ) 33 , b ) 50 , c ) 66 , d ) 10 , e ) 18\n",
      "Let's think step by step\n",
      "Michael started with 58 golf balls.\n",
      "After losing 23 on tues- day, he had 58 - 23 = 35.\n",
      "After losing 2 more, he had 35 - 2 = 33 golf balls.\n",
      "The answer is a, 33.\n",
      "\n",
      "Question: Olivia has $23. She bought five bagels for $3 each. How much money does she have left?\n",
      "Let's think step by step\n",
      "Options: a ) 3 , b ) 5 , c ) 8 , d ) 10 , e ) none of these\n",
      "Olivia had 23 dollars.\n",
      "5 bagels for 3 dollars each will be 5 x 3 = 15 dollars.\n",
      "So she has 23 - 15 dollars left.\n",
      "23 - 15 is 8.\n",
      "The answer is c, 8.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(prompt_original)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "083f2a1c-6a3b-41f5-a85b-2512a043a2d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_complex = open('MathQA/prompt_complex.txt').read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "7d417ccb-bc73-4eef-b07a-447606170a07",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Question: there were 35 students in a hostel . due to the admission of 7 new students the expenses of the mess were increased by rs . 84 per day while the average expenditure per head diminished by re 1 . what was the original expenditure of the mess ?\n",
      "Options: a ) rs 450 , b ) rs 920 , c ) rs 550 , d ) rs . 630 , e ) none of these\n",
      "Let's think step by step\n",
      "let the original average expenditure be rs . x\n",
      "then , 42 ( x - 1 ) - 35 x = 84\\,7 x = 126\n",
      "x = 18\n",
      "therefore original expenditure = rs . ( 35 * 18 ) = rs . 630.\n",
      "The answer is d, rs . 630\n",
      "\n",
      "Question: a train 200 m long passes a man , running at 5 km / hr in the same direction in which the train is going , in 10 seconds . the speed of the train is ?\n",
      "Options: a ) 28 , b ) 50 , c ) 77 , d ) 22 , e ) 12\n",
      "Let's think step by step\n",
      "speed of the train relative to man = ( 200 / 10 ) m / sec = ( 20 ) m / sec.\n",
      "[ ( 20 ) * ( 18 / 5 ) ] km / hr = 72 km / hr.\n",
      "let the speed of the train be x km / hr.\n",
      "then , relative speed = ( x - 5 ) km / hr.\n",
      "x - 5 = 72, x = 77 km / hr .\n",
      "The answer is c, 77\n",
      "\n",
      "Question: solution x contains 20 % of material a and 80 % of material b . solution y contains 30 % of material a and 70 % of material b . a mixture of both these solutions contains 22 % of material a in the final product . how much solution x is present in the mixture ?\n",
      "Options: a ) 40 % , b ) 60 % , c ) 80 % , d ) 100 % , e ) 110 %\n",
      "Let's think step by step\n",
      "we can assume the total weight of the mixture = 100\n",
      "conc of a in the final mixture = 22\n",
      "let weight of a in the mixture be x.\n",
      "conc given = 20% = 0.2\n",
      "therefore , weight of b = 100 - x.\n",
      "conc given = 30% = 0.3\n",
      "now , accordding to the problem, 0.2 x + 0.3 ( 100 - x ) = 22\n",
      "solving , we get x = 80\n",
      "since we assumed the weight of the mixture = 100, therefore presence of a in the mixture = 80%.\n",
      "The answer is c, 80%\n",
      "\n",
      "Question: a trader sells 40 metres of cloth for rs . 8200 at a profit of rs . 35 per metre of cloth . how much profit will the trder earn on 40 metres of cloth ?\n",
      "Options: a ) rs . 950 , b ) rs . 1500 , c ) rs . 1000 , d ) rs . 1400 , e ) none of these\n",
      "Let's think step by step\n",
      "price of 1 metre cloth = 8200 / 40 = rs 205\n",
      "cost of 1 metre cloth = rs 205 – 35 = rs 170\n",
      "cost on 40 metres = 170 x 40 = rs . 6800\n",
      "profit earned on 40 metres cloth = rs . 8200 – rs . 6800 = rs . 1400\n",
      "The answer is d, rs . 1400\n",
      "\n",
      "Question: if x < y < z and y - x > 5 , where x is an even integer and y and z are odd integers , what is the least possible value s of z - x ?\n",
      "Options: a ) 6 , b ) 7 , c ) 8 , d ) 9 , e ) 10\n",
      "Let's think step by step\n",
      "We know x < y < z\n",
      "to find the least possible value for z - x, we need to find the values for z and x that can be closest to each other.\n",
      "if x is some even number, then what could be minimum possible odd z.\n",
      "if x is some even number, y - x > 5 ; y > x + 5\n",
      "minimum value for y = x + 5 + 2 = x + 7\n",
      "(note : x + 5 is as even + odd = odd and nearest odd greater than x + 5 is x + 5 + 2)\n",
      "minimum value for z = y + 2 = x + 7 + 2 = x + 9\n",
      "(note : z = y + 2 because both z and y are odd, difference between two odd numbers is 2)\n",
      "s = z - x = x + 9 - x = 9\n",
      "The answer is d, 9\n",
      "\n",
      "Question: what is the difference between the c . i . on rs . 6000 for 1 1 / 2 years at 4 % per annum compounded yearly and half - yearly ?\n",
      "Options: a ) s . 2.04 , b ) s . 2.08 , c ) s . 2.02 , d ) s . 2.83 , e ) s . 2.45\n",
      "Let's think step by step\n",
      "c . i . when interest is compounded yearly = [ 6000 * ( 1 + 4 / 100 ) * ( 1 + ( 1 / 2 * 4 ) / 100 ] = 6000 * 26 / 25 * 51 / 50 = rs . 6364.8\n",
      "c . i . when interest is compounded half - yearly = [ 6000 * ( 1 + 2 / 100 ) 2 ] = ( 6000 * 51 / 50 * 51 / 50 * 51 / 50 ) = rs . 6367.25\n",
      "difference = ( 6367.25 - 6364.8 ) = rs . 2.45.\n",
      "The answer is e, s . 2.45\n",
      "\n",
      "Question: the average weight of a , b and c is 45 kg . if the average weight of a and b be 40 kg and that of b and c be 45 kg , then the weight of b is :\n",
      "Options: a ) 31 kg , b ) 32 kg , c ) 33 kg , d ) 35 kg , e ) none of theseLet's think step by step\n",
      "let a , b , c represent their respective weights.\n",
      "then , we have : a + b + c = ( 45 x 3 ) = 135 … ( i )\n",
      "a + b = ( 40 x 2 ) = 80 … ( ii )\n",
      "b + c = ( 45 x 2 ) = 90 … ( iii )\n",
      "adding ( ii ) and ( iii ), we get : a + 2 b + c = 170 … ( iv )\n",
      "subtracting ( i ) from ( iv ), we get : b = 35.\n",
      "b’ s weight = 35 kg .\n",
      "The answer is d 35 kg\n",
      "\n",
      "Question: the compound and the simple interests on a certain sum at the same rate of interest for two years are rs . 11730 and rs . 10200 respectively . find the sum\n",
      "Options: a ) rs . 17037 , b ) rs . 17000 , c ) rs . 17276 , d ) rs . 170287 , e ) rs . 171881\n",
      "Let's think step by step\n",
      "the simple interest for the first year is 10200 / 2 is rs . 5100\n",
      "and compound interest for first year also is rs . 5100.\n",
      "the compound interest for second year on rs . 5100 for one year\n",
      "so rate of the interest = ( 100 * 1530 ) / ( 5100 * 1 ) = 30% p\n",
      "so p = ( 100 * 10200 ) / ( 30 * 2 ) = rs . 17000\n",
      "The answer is b, rs . 17000\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(prompt_complex)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "3f137d33-cace-4fbe-8de6-f7a4807a1df8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_answer(pred_str, ans_str):\n",
    "    \"\"\"Find the last number as the predicted answer\"\"\"\n",
    "    pattern = '\\d*\\.?\\d+'\n",
    "    pred = re.findall(pattern, pred_str)\n",
    "    if(len(pred) >= 1):\n",
    "        # print(pred_str)\n",
    "        pred = pred[-1]\n",
    "        gold = re.findall(pattern, ans_str)\n",
    "        # print(ans_str)\n",
    "        gold = gold[-1]\n",
    "        return pred == gold\n",
    "    else: return False\n",
    "\n",
    "def parse_pred_ans(filename):\n",
    "    with open(filename) as fd: lines = fd.readlines()\n",
    "    am, a = None, None\n",
    "    num_q, acc = 0, 0\n",
    "    current_mode = 'none'\n",
    "    questions = []\n",
    "    ans_pred = []\n",
    "    ans_gold = []\n",
    "    for l in lines:\n",
    "        if(l.startswith('Question: ')):\n",
    "            if(am is not None and a is not None):\n",
    "                questions.append(q)\n",
    "                ans_pred.append(am)\n",
    "                ans_gold.append(a)\n",
    "                if(test_answer(am, a)):\n",
    "                    acc += 1\n",
    "            current_mode = 'q'\n",
    "            q = l\n",
    "            num_q += 1\n",
    "        elif(l.startswith('Options: ')): pass\n",
    "        elif(l.startswith('A_model:')):\n",
    "            current_mode = 'am'\n",
    "            am = l\n",
    "        elif(l.startswith('A:')):\n",
    "            current_mode = 'a'\n",
    "            a = l\n",
    "        else:\n",
    "            if(current_mode == 'q'): q += l\n",
    "            elif(current_mode == 'am'): am += l\n",
    "            elif(current_mode == 'a'): a += l\n",
    "            else:\n",
    "                raise ValueError(current_mode)\n",
    "                \n",
    "    questions.append(q)\n",
    "    ans_pred.append(am)\n",
    "    ans_gold.append(a)\n",
    "    if(test_answer(am, a)):\n",
    "        acc += 1\n",
    "    print('num_q %d correct %d ratio %.4f' % (num_q, acc, float(acc / num_q)))\n",
    "    return questions, ans_pred, ans_gold"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b0456b8-89d6-4058-8aa5-56ca6142d4d2",
   "metadata": {},
   "source": [
    "# Original Prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "d2ad4aef-22b8-4e95-a9c8-9dc87e582784",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_ans_dict(c):\n",
    "    if(c.startswith(\"['\")):\n",
    "        c_ = eval(c)\n",
    "        ans_dict = {}\n",
    "        for ac in c_:\n",
    "            idx_, num_ = ac.split(' ) ')[-2:]\n",
    "            if('a' in idx_): ans_dict['a'] = num_\n",
    "            if('b' in idx_): ans_dict['b'] = num_\n",
    "            if('c' in idx_): ans_dict['c'] = num_\n",
    "            if('d' in idx_): ans_dict['d'] = num_\n",
    "            if('e' in idx_): ans_dict['e'] = num_\n",
    "    else:\n",
    "        ans_dict = {}\n",
    "        pattern = ' , [a-e] \\) '\n",
    "        c_ = re.split(pattern, c)\n",
    "        c_[0] = c_[0].split(' ) ')[-1]\n",
    "        for i, c in zip(['a', 'b', 'c', 'd', 'e'], c_):\n",
    "            ans_dict[i] = c\n",
    "    return ans_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "9aa692a1-b645-4149-ade2-c6ee83964601",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 76%|██████████████████████████████████████████████████████████████████████████████████████████████████▊                               | 456/600 [31:23<09:54,  4.13s/it]\n"
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "not enough values to unpack (expected 2, got 1)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Input \u001b[0;32mIn [24]\u001b[0m, in \u001b[0;36m<cell line: 2>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      9\u001b[0m ans_dict \u001b[38;5;241m=\u001b[39m {}\n\u001b[1;32m     10\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m ac \u001b[38;5;129;01min\u001b[39;00m c\u001b[38;5;241m.\u001b[39msplit(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m, \u001b[39m\u001b[38;5;124m'\u001b[39m):\n\u001b[0;32m---> 11\u001b[0m     idx_, num_ \u001b[38;5;241m=\u001b[39m ac\u001b[38;5;241m.\u001b[39msplit(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m ) \u001b[39m\u001b[38;5;124m'\u001b[39m)[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m2\u001b[39m:]\n\u001b[1;32m     12\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124ma\u001b[39m\u001b[38;5;124m'\u001b[39m \u001b[38;5;129;01min\u001b[39;00m idx_): ans_dict[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124ma\u001b[39m\u001b[38;5;124m'\u001b[39m] \u001b[38;5;241m=\u001b[39m num_\n\u001b[1;32m     13\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mb\u001b[39m\u001b[38;5;124m'\u001b[39m \u001b[38;5;129;01min\u001b[39;00m idx_): ans_dict[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mb\u001b[39m\u001b[38;5;124m'\u001b[39m] \u001b[38;5;241m=\u001b[39m num_\n",
      "\u001b[0;31mValueError\u001b[0m: not enough values to unpack (expected 2, got 1)"
     ]
    }
   ],
   "source": [
    "i = 0\n",
    "with open('outputs/test_outputs_original.txt', 'w') as fd:\n",
    "    for case in tqdm(test_data):\n",
    "        q = case['Problem']\n",
    "        c = case['options']\n",
    "        a = case['Rationale']\n",
    "        cr = case['correct']\n",
    "    \n",
    "        ans_dict = {}\n",
    "        for ac in c.split(', '):\n",
    "            idx_, num_ = ac.split(' ) ')[-2:]\n",
    "            if('a' in idx_): ans_dict['a'] = num_\n",
    "            if('b' in idx_): ans_dict['b'] = num_\n",
    "            if('c' in idx_): ans_dict['c'] = num_\n",
    "            if('d' in idx_): ans_dict['d'] = num_\n",
    "            if('e' in idx_): ans_dict['e'] = num_\n",
    "            \n",
    "        a += ' ' + ans_dict[cr]\n",
    "        \n",
    "        prompt_q = '' \n",
    "        prompt_q += 'Question: ' + q + '\\n'\n",
    "        prompt_q += 'Options: ' + c + '\\n'\n",
    "        prompt_q += \"Let's think step by step\\n\"\n",
    "        \n",
    "        prompt_full = ''.join(prompt_original) + prompt_q\n",
    "        response = openai.Completion.create(model=\"text-davinci-002\", \n",
    "                                            prompt=prompt_full, \n",
    "                                            temperature=0, \n",
    "                                            max_tokens=256)\n",
    "        \n",
    "        ans_model = response['choices'][0]['text']\n",
    "        fd.write('%sA_model:\\n%s\\nA:\\n%s\\n\\n' % (prompt_q, ans_model, a))\n",
    "        i += 1\n",
    "        # if(i == 2): break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "27b3d137-5abe-4f21-a414-8fb56874cd89",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 600/600 [09:50<00:00,  1.02it/s]\n"
     ]
    }
   ],
   "source": [
    "i = 0\n",
    "with open('outputs/test_outputs_original_from_456.txt', 'w') as fd:\n",
    "    for case in tqdm(test_data):\n",
    "        if(i < 456):\n",
    "            i += 1\n",
    "            continue\n",
    "        \n",
    "        q = case['Problem']\n",
    "        c = case['options']\n",
    "        a = case['Rationale']\n",
    "        cr = case['correct']\n",
    "    \n",
    "        ans_dict = get_ans_dict(c)\n",
    "            \n",
    "        a += ' ' + ans_dict[cr]\n",
    "        \n",
    "        prompt_q = '' \n",
    "        prompt_q += 'Question: ' + q + '\\n'\n",
    "        prompt_q += 'Options: ' + c + '\\n'\n",
    "        prompt_q += \"Let's think step by step\\n\"\n",
    "        \n",
    "        prompt_full = ''.join(prompt_original) + prompt_q\n",
    "        response = openai.Completion.create(model=\"text-davinci-002\", \n",
    "                                            prompt=prompt_full, \n",
    "                                            temperature=0, \n",
    "                                            max_tokens=256)\n",
    "        \n",
    "        ans_model = response['choices'][0]['text']\n",
    "        fd.write('%sA_model:\\n%s\\nA:\\n%s\\n\\n' % (prompt_q, ans_model, a))\n",
    "        i += 1\n",
    "        # if(i == 2): break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "c2f1ccfd-6662-42ef-91de-61d001c0eadd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "num_q 456 correct 133 ratio 0.2917\n"
     ]
    }
   ],
   "source": [
    "_, _, _ = parse_pred_ans('outputs/test_outputs_original.txt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "92479717-2c3a-4de2-9cd8-09926151b936",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "num_q 145 correct 48 ratio 0.3310\n"
     ]
    }
   ],
   "source": [
    "_, _, _ = parse_pred_ans('outputs/test_outputs_original_from_456.txt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "35a569ce-0b53-4aa6-8752-571cca043f3f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.3011647254575707"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(133 + 48) / (456 + 145)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77a747e6-7934-47b6-af0e-a816a2ebf4b1",
   "metadata": {},
   "source": [
    "# Complex Prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "e5e08bf6-444e-427d-ab14-21aa51670089",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 600/600 [43:57<00:00,  4.40s/it]\n"
     ]
    }
   ],
   "source": [
    "i = 0\n",
    "with open('outputs/test_outputs_complex.txt', 'w') as fd:\n",
    "    for case in tqdm(test_data):\n",
    "        q = case['Problem']\n",
    "        c = case['options']\n",
    "        a = case['Rationale']\n",
    "        cr = case['correct']\n",
    "    \n",
    "        ans_dict = get_ans_dict(c)\n",
    "            \n",
    "        a += ' ' + ans_dict[cr]\n",
    "        \n",
    "        prompt_q = '' \n",
    "        prompt_q += 'Question: ' + q + '\\n'\n",
    "        prompt_q += 'Options: ' + c + '\\n'\n",
    "        prompt_q += \"Let's think step by step\\n\"\n",
    "        \n",
    "        prompt_full = ''.join(prompt_complex) + prompt_q\n",
    "        response = openai.Completion.create(model=\"text-davinci-002\", \n",
    "                                            prompt=prompt_full, \n",
    "                                            temperature=0, \n",
    "                                            max_tokens=256)\n",
    "        \n",
    "        ans_model = response['choices'][0]['text']\n",
    "        fd.write('%sA_model:\\n%s\\nA:\\n%s\\n\\n' % (prompt_q, ans_model, a))\n",
    "        i += 1\n",
    "        # if(i == 2): break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "c8e9cc8e-7682-4bc5-9e5c-270dce937145",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "num_q 600 correct 216 ratio 0.3600\n"
     ]
    }
   ],
   "source": [
    "_, _, _ =parse_pred_ans('outputs/test_outputs_complex.txt')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9334885-0c97-4b61-88fe-976c6df44791",
   "metadata": {},
   "source": [
    "# Random Prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "ee89e639-967d-43c2-8189-bc7574944444",
   "metadata": {},
   "outputs": [],
   "source": [
    "random_idx = np.random.choice(len(train_data), 8, replace=False)\n",
    "np.save('MathQA/random_prompt_idx.npy', random_idx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "86d0c162-269d-45db-b84e-2afe3a014564",
   "metadata": {},
   "outputs": [],
   "source": [
    "random_idx = np.load('MathQA/random_prompt_idx.npy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "d22bc677-7d78-49f2-8fc4-056aa88943a9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([20632, 15959, 12471, 21233, 17351, 14350, 19776,  8829])"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "random_idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "5117d4a7-f041-498b-8aaa-ab66ac2ffe00",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'Problem': 'of the two - digit positive integers that have no digits equal to zero , how many have two digits that are equal to each other and the remaining digit different from the other two ?',\n",
       " 'Rationale': 'of the two - digit positive integers that have no digits equal to zero , how many have two digits that are equal to each other and the remaining digit different from the other two ? a . 24 b . 36 c . 72 d . 144 e . 216 choosing the digit for x - 9 ways ; # of permutations of 3 digits in xx - 2 ! / 2 ! total : 9 * 2 ! / 2 ! = 9 . answer : b .',\n",
       " 'options': 'a ) 24 , b ) 9 , c ) 72 , d ) 144 , e ) 216',\n",
       " 'correct': 'b',\n",
       " 'annotated_formula': 'multiply(multiply(subtract(const_10, const_1), subtract(const_10, const_2)), const_1)',\n",
       " 'linear_formula': 'subtract(const_10,const_1)|subtract(const_10,const_2)|multiply(#0,#1)|multiply(#2,const_1)',\n",
       " 'category': 'general'}"
      ]
     },
     "execution_count": 61,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_data[20632]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "1a093b4d-46ee-49ef-ac90-bae16e2fe192",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('MathQA/prompt_random.txt', 'w') as fd:\n",
    "    for i in random_idx:\n",
    "        case = train_data[i]\n",
    "        fd.write('Question: ' + case['Problem'] + '\\n')\n",
    "        fd.write('Options: ' + case['options'] + '\\n')\n",
    "        fd.write(\"Let's think step by step\\n\")\n",
    "        fd.write(case['Rationale'] + '\\n')\n",
    "        fd.write('The answer is ' + case['correct'] + '\\n\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "0ce134a3-d22a-4020-ba44-c16621eb4c7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_random = open('MathQA/prompt_random.txt').read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "4f03f179-deb5-4e82-9abc-326087ffece0",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 600/600 [49:47<00:00,  4.98s/it]\n"
     ]
    }
   ],
   "source": [
    "i = 0\n",
    "with open('outputs/test_outputs_random.txt', 'w') as fd:\n",
    "    for case in tqdm(test_data):\n",
    "        q = case['Problem']\n",
    "        c = case['options']\n",
    "        a = case['Rationale']\n",
    "        cr = case['correct']\n",
    "    \n",
    "        ans_dict = get_ans_dict(c)\n",
    "            \n",
    "        a += ' ' + ans_dict[cr]\n",
    "        \n",
    "        prompt_q = '' \n",
    "        prompt_q += 'Question: ' + q + '\\n'\n",
    "        prompt_q += 'Options: ' + c + '\\n'\n",
    "        prompt_q += \"Let's think step by step\\n\"\n",
    "        \n",
    "        prompt_full = ''.join(prompt_random) + prompt_q\n",
    "        response = openai.Completion.create(model=\"text-davinci-002\", \n",
    "                                            prompt=prompt_full, \n",
    "                                            temperature=0, \n",
    "                                            max_tokens=256)\n",
    "        \n",
    "        ans_model = response['choices'][0]['text']\n",
    "        fd.write('%sA_model:\\n%s\\nA:\\n%s\\n\\n' % (prompt_q, ans_model, a))\n",
    "        i += 1\n",
    "        # if(i == 2): break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "d6366b4b-98a9-4566-9e4e-c7061b11ab92",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "num_q 600 correct 209 ratio 0.3483\n"
     ]
    }
   ],
   "source": [
    "_, _, _ =parse_pred_ans('outputs/test_outputs_random.txt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86645c81-9b61-481c-9542-9a7c35e9c660",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
